--- Functors representing data structures that can be traversed from left to right.
package Data.Traversable where

import frege.Prelude hiding (mapM, forM, for, sequence, foldr)
import Data.Foldable (Foldable)
import Data.Monoid
import frege.data.wrapper.Identity
import frege.data.wrapper.Const

{-- 
   Functors representing data structures that can be traversed from left to right.

   Minimal complete definition: 'traverse' or 'sequenceA'.
 
   The superclass instances should satisfy the following:

    * In the 'Functor' instance, 'fmap' should be equivalent to traversal
      with the identity applicative functor ('fmapDefault').

    * In the 'Foldable' instance, 'Foldable.foldMap' should be
      equivalent to traversal with a constant applicative functor
      ('foldMapDefault').
      
    Note that the functions 'mapM', 'sequence', 'forM' are just specialized versions 
    of 'traverse', 'sequenceA' and 'for', and wouldn't be required in Frege.
    They are included for Haskell compatibility only. In Haskell the specialized 
    functions were needed because Haskell Monads haven't been Applicatives before GHC 7.10.
   -}
class (Functor t, Foldable t) => Traversable t where
    {-- Map each element of a structure to an action, evaluate
        these actions from left to right, and collect the results.
     -}
    traverse :: Applicative f => (a -> f b) -> t a -> f (t b)
    traverse f ts = sequenceA $ fmap f ts

    {-- Evaluate each action in the structure from left to right,
        and collect the results.
    -}
    sequenceA :: Applicative f => t (f a) -> f (t a)
    sequenceA ts = traverse id ts

    {-- Map each element of a structure to a monadic action, evaluate
        these actions from left to right, and collect the results.
        This function exists for Haskell compatibility only.
    -}
    mapM :: Monad m => (a -> m b) -> t a -> m (t b)
    mapM f ts = traverse f ts

    {-- Evaluate each monadic action in the structure from left to right,
        and collect the results.
        This function exists for Haskell compatibility only.
    -}
    sequence :: Monad m => t (m a) -> m (t a)
    sequence ts = mapM id ts

-- instances for Prelude types

instance Traversable Maybe where
    traverse _ Nothing = pure Nothing
    traverse f (Just x) = fmap Just (f x)

instance Traversable [] where
    traverse f ts = foldr cons_f (pure []) ts
      where cons_f x ys = fmap (:) (f x) <*> ys
    mapM = Prelude.mapM
      
instance Traversable Identity where
    traverse f (Identity x) = Identity `fmap` f x

-- general functions

--- 'for' is 'traverse' with its arguments flipped.
for :: (Traversable t, Applicative f) => t a -> (a -> f b) -> f (t b)
for ts f = traverse f ts

--- 'forM' is 'mapM' with its arguments flipped.
---  This function exists for Haskell compatibility only.
forM :: (Traversable t, Monad m) => t a -> (a -> m b) -> m (t b)
forM ts f = mapM f ts

-- left-to-right state transformer
-- TODO: make private when bug is fixed
protected data StateL s a = StateL { run :: s -> (s, a) }

protected instance Functor (StateL s) where
    fmap f (StateL.StateL k) = StateL.StateL $ (\ s -> let (s', v) = k s in (s', f v))

protected instance Applicative (StateL s) where
    pure x = StateL.StateL (\ s -> (s, x))
    StateL.StateL kf <*> StateL.StateL kv = StateL.StateL (\ s ->
        let (s', f) = kf s
            (s'', v) = kv s'
        in (s'', f v))

{-- The 'mapAccumL' function behaves like a combination of 'fmap'
    and 'fold'; it applies a function to each element of a structure,
    passing an accumulating parameter from left to right, and returning
    a final value of this accumulator together with the new structure.
    -}
mapAccumL :: Traversable t => (a -> b -> (a, c)) -> a -> t b -> (a, t c)
mapAccumL f s t = StateL.run (traverse (StateL.StateL . flip f) t) s

-- right-to-left state transformer
-- TODO: make private when bug is fixed
protected data StateR s a = StateR { run :: s -> (s, a) }

protected instance Functor (StateR s) where
    fmap f (StateR.StateR k) = StateR.StateR (\ s -> let (s', v) = k s in (s', f v))

protected instance Applicative (StateR s) where
    pure x = StateR.StateR (\ s -> (s, x))
    StateR.StateR kf <*> StateR.StateR kv = StateR.StateR (\ s ->
        let (s', v) = kv s
            (s'', f) = kf s'
        in (s'', f v))

{-- The 'mapAccumR' function behaves like a combination of 'fmap'
    and 'foldr'; it applies a function to each element of a structure,
    passing an accumulating parameter from right to left, and returning
    a final value of this accumulator together with the new structure.
    -}
mapAccumR :: Traversable t => (a -> b -> (a, c)) -> a -> t b -> (a, t c)
mapAccumR f s t = StateR.run (traverse (StateR.StateR . flip f) t) s

{-- This function may be used as a value for `fmap` in a `Functor`
    instance, provided that 'traverse' is defined. (Using
    'fmapDefault' with a 'Traversable' instance defined only by
    'sequenceA' will result in infinite recursion.)
    -}
fmapDefault :: Traversable t => (a -> b) -> t a -> t b
fmapDefault f ts = Identity.run $ traverse (Identity . f) ts

{-- This function may be used as a value for `Data.Foldable.foldMap`
    in a `Foldable` instance.
-}
foldMapDefault :: (Traversable t, Monoid m) => (a -> m) -> t a -> m
foldMapDefault f ts = Const.get $ traverse (Const . f) ts
